import segyio
import matplotlib.pyplot as plt
import numpy as np
import argparse

# 设置程序参数
parser = argparse.ArgumentParser()
parser.add_argument("file_name", type=str, help="name of the segy file")
parser.add_argument('--gain_type', type=int, default=1, help="gain every trace: 1, gain all trace: 2")
parser.add_argument("--percent", type=float, default=95, help="gain percent")
opt = parser.parse_args()
print(opt)

# 下一炮读取的起始道下标
start_trace = 0

# 起始道列表, 用于记录已经加载过炮集的起始道
loaded_start_traces = []

# 增益类型(逐道增益=1, 全炮增益=2)
gain_type = opt.gain_type

# 增益百分比
percent = opt.percent

# segy文件名
segy_filename = opt.file_name

# 打开segy文件
segy_file = None
try:
    segy_file = segyio.open(segy_filename, ignore_geometry=True)
except Exception as e:
    print(e)
    print('Load file "%s" Failed!' % segy_filename)
    exit(-1)


# 从start_trace开始读取一炮数据
def load_trace_set(start_trace):
    if start_trace >= segy_file.tracecount:
        return 0, None, start_trace

    rec_num = segy_file.header[start_trace][segyio.TraceField.FieldRecord]
    end_trace = segy_file.tracecount
    for i in range(start_trace+1, segy_file.tracecount):
        if segy_file.header[i][segyio.TraceField.FieldRecord] != rec_num:
            end_trace = i
            break

    print("rec_num: %d, %d -> %d" % (rec_num, start_trace+1, end_trace))
    m = segy_file.trace.raw[start_trace: end_trace]

    if gain_type == 1:
        per_max = np.percentile(np.abs(m), 100, 1)
        m = np.array(list(map(lambda x, y: x / y, m, per_max)))
    else:
        per_max = np.percentile(np.abs(m), percent)
        m = m / per_max

    return rec_num, m.T, end_trace


# 显示地震剖面
def show_graph(rec_num, m):
    plt.cla()
    plt.suptitle('Filed Record Number: ' + str(rec_num), fontsize=16)
    plt.imshow(m, cmap=plt.cm.seismic, vmin=-1, vmax=1, interpolation='bilinear')


# 键盘事件处理函数
def on_key(event):
    global start_trace, fig
    if event.key == 'right':
        # 载入下一炮数据
        prev_start = start_trace
        rec_num, m, start_trace = load_trace_set(start_trace)
        if rec_num <= 0:
            return
        loaded_start_traces.append(prev_start)
        # 显示炮集数据
        show_graph(rec_num, m)
        plt.draw()
    elif event.key == 'left':
        if len(loaded_start_traces) > 1:
            # 载入上一炮数据
            loaded_start_traces.pop()
            rec_num, m, start_trace = load_trace_set(loaded_start_traces[-1])
            # 显示炮集数据
            show_graph(rec_num, m)
            plt.draw()
    elif event.key == 'escape':
        # 退出显示程序
        plt.close(fig)
        return


if __name__ == '__main__':
    # 设置plt显示对象
    fig = plt.figure()
    fig.canvas.mpl_connect('key_press_event', on_key)
    ax = fig.add_subplot(111)
    ax.set_aspect(20)

    # 加载第一炮数据
    rec_num, m, start_trace = load_trace_set(start_trace)

    # 如果加载成功打开绘图窗口
    if m is not None:
        loaded_start_traces.append(0)
        show_graph(rec_num, m)
        plt.show()

segy_file.close()
